import os
import time

from tqdm import tqdm
import google.generativeai as genai
import PIL.Image

## personal api key here
genai.configure(api_key = '#UPDATE#HERE#')

model = genai.GenerativeModel('gemini-pro-vision')

files = os.listdir("/Users/yuwang/Desktop/Gemini in Political Science/data/images")

eval_errors = 0
key_errors = 0
max_tries = 5

def image_analysis(image_name):
    print("Image-name：", image_name)
    img = PIL.Image.open("./../data/images/" + image_name)
    response = model.generate_content(["List all the objects and people in the image, using lowercase and singular forms. Include the count for each object. Please use the following one-line format with curly brackets in the response: {\"Object 1\": Count 1, \"Object 2\": Count 2}. If there are no objects, return None.", img])
    print(response.text)
    if response.text.strip() == "None":
        return response.text.strip()
    try:
        result = eval(response.text)
    except:
        eval_errors += 1
    keys = result.keys()
    for key in keys:
        key = key.replace(" ", "").replace("-", "")
        if not key.isalpha():
            key_errors += 1
            raise Exception("Identified objects contains non-letter characters.")
    return response.text

def main():
    with open("results.txt", "w") as doc:
        for file in tqdm(files):
            n_tries = 0
            success = False
            while n_tries < max_tries and not success:
                try:
                    n_tries += 1
                    result = image_analysis(file)
                    doc.write(file + "\t" + result + "\n")
                    success = True
                except:
                    print(file, n_tries)
            if not success:
                doc.write(file + "\t" + "Error" + "\n")
if __name__=="__main__": 
    start = time.time()
    main() 
    print("Amount of time", time.time() - start)
